/*
  non-Newtonian fluid with pressure- and shear-dependent viscosity
  
  Viscosity is of the form
   
     mu_0(A+|Du|^2)^((r-2)/2).
   
*/
#include <math.h>
#include <petscksp.h>
#include "fem.h"
#include "cell_locator.h"
#include "tof.h"


#define GAP 0.1

class C3DNonNewtonianStokesProblem : public CStateProblem
{
 private:
  int    conv;
  double PAR_A, PAR_mu0, PAR_r, PAR_sigma1, PAR_sigma2, delta;
  CVar   diverg, visco;
  map<int, bool> outflows; // boundaries where outflow b.c. is applied
  map<int, double> outflow_p; // outflow/inflow pressures
  int mortar_pid, mortar_sid; // primary and secondary surface id for mortaring
  CellLocator *locator;       // face locator for mortaring
 
 public:
  ~C3DNonNewtonianStokesProblem();
  CVec &getQ() { return q; };
  
  void afterReadConfig();
  void afterPrepare();
  void afterEvaluateVolumeComponents(int elem, int qp);
  
  bool useVolumeIntegral() { return true; };
  bool useFaceIntegral() { return true; };
  bool useEdgeIntegral() { return false; };
  CVar volumeIntegral(int comp, int elem, int qp);
  CVar faceIntegral(int comp, int elem);
  CVar edgeIntegral(int comp, int elem);
  
  void saveOutflow();
  void savePressure();
};


C3DNonNewtonianStokesProblem::~C3DNonNewtonianStokesProblem()
{
  outflows.clear();
  outflow_p.clear();
}

void C3DNonNewtonianStokesProblem::afterReadConfig()
{
  bool   setPressureMeanValue;
  int    pressureMVsegment = 0;
  double pressureMV = 0;
  
  if (params.find("pressure_meanvalue") != params.end())
  {
    setPressureMeanValue = true;
    int comma = params["pressure_meanvalue"].find(",");
    string segment = params["pressure_meanvalue"].substr(0,comma),
           mvalue  = params["pressure_meanvalue"].substr(comma+1);
	   
    pressureMVsegment = atoi(segment.c_str());
    pressureMV = atof(mvalue.c_str());
    log(W_PROG, "Pressure meanvalue on %s = %e\n", mesh->getPhysicalName(pressureMVsegment).c_str(), pressureMV);
  }
  else
  {
    setPressureMeanValue = false;
  }
  
  if (params.find("mortar_interfaces") != params.end())
  {
	  int comma = params["mortar_interfaces"].find(",");
	  string primary = params["mortar_interfaces"].substr(0,comma),
			 secondary = params["mortar_interfaces"].substr(comma+1);

	  mortar_pid = mesh->getPhysNameId(primary);
	  mortar_sid = mesh->getPhysNameId(secondary);

	  locator = new CellLocator(mesh, mortar_sid, 10, 10, 1);
  }

  registerQuadratures(&Q3d11p, &Q2d6p, &Q1d4p);
  addVector("velocity", FE_LIN);
  addScalar("pressure", FE_LIN, setPressureMeanValue, pressureMVsegment, pressureMV);
  
//  C3DFlowProblem::afterReadConfig();
}

void C3DNonNewtonianStokesProblem::afterPrepare()
{
	char e[1023];

	// start with initial guess different from zero
	if (!loaded_solution)
	{
		for (int i=0; i<q.size(); q(i) = q(i) + 1e-0, i++);
	}
	else
	{
	//    for (int i=0; i<q.size(); q(i) = q(i) - q(i).getVal()*0.9, i++);
	}

	if (params.find("outflow") != params.end())
	{
		size_t begin = 0;
		size_t comma = 0;
		string bname;

		while (comma < string::npos)
		{
		  comma = params["outflow"].find_first_of(",", begin);
		  bname = params["outflow"].substr(begin, comma-begin);
		  outflows[mesh->getPhysNameId(bname)] = true;
		  begin = comma+1;

		  if (params.find("outflow_pressure[" + bname + "]") != params.end())
			outflow_p[mesh->getPhysNameId(bname)] = atof(params["outflow_pressure[" + bname + "]"].c_str());
		}
	}

	// parameter convection indicates whether convective term is to be included into the equations
	if (params.find("convection") != params.end() && params["convection"].compare("yes") == 0)
	{
		conv = 1;
		log(W_PROG, "\nUsing convective term.\n");
	}
	else
	{
		conv = 0;
		log(W_PROG, "\nNot using convective term.\n");
	}

	// model of viscosity
	PAR_mu0 = atof(params["parameter_mu0"].c_str());
	PAR_A     = atof(params["parameter_A"].c_str());
	PAR_r     = atof(params["parameter_r"].c_str());
	log(W_PROG, "\nViscosity parameters:\n");
	log(W_PROG, "A     = %e\nr     = %e\nmu0   = %e\n\n", PAR_A, PAR_r, PAR_mu0);
	if (PAR_r != 2)
	{
	  sprintf(e, "%e*(%e+gvelocity00^2+gvelocity11^2+gvelocity22^2+(gvelocity01+gvelocity10)^2*0.5+(gvelocity02+gvelocity20)^2*0.5+(gvelocity12+gvelocity21)^2*0.5)^(%e)", PAR_mu0, PAR_A, (PAR_r-2)/2);
	  addPPscalar("viscosity", e);
	}
	else
	{
	  sprintf(e, "%e", PAR_mu0);
	  addPPscalar("viscosity", e);
	}

	// coefficients of outflow boundary conditions
	PAR_sigma1= atof(params["parameter_sigma1"].c_str());
	PAR_sigma2= atof(params["parameter_sigma2"].c_str());
	log(W_PROG, "Outflow coefficients:\nsigma1 = %e\nsigma2 = %e\n\n", PAR_sigma1, PAR_sigma2);

	addPPscalar("normDv", "sqrt(gvelocity00^2+gvelocity11^2+gvelocity22^2+(gvelocity01+gvelocity10)^2*0.5+(gvelocity02+gvelocity20)^2*0.5+(gvelocity12+gvelocity21)^2*0.5)");
}

void C3DNonNewtonianStokesProblem::afterEvaluateVolumeComponents(int elem, int qp)
{
  if (PAR_r != 2)
  {
	CVar Du2 = U[0]->grad(0)*U[0]->grad(0)
		+U[1]->grad(1)*U[1]->grad(1)
		+U[2]->grad(2)*U[2]->grad(2)
		+(U[0]->grad(1)+U[1]->grad(0))*(U[0]->grad(1)+U[1]->grad(0))*0.5
		+(U[0]->grad(2)+U[2]->grad(0))*(U[0]->grad(2)+U[2]->grad(0))*0.5
		+(U[1]->grad(2)+U[2]->grad(1))*(U[1]->grad(2)+U[2]->grad(1))*0.5;
	visco  = pow(Du2 + PAR_A, (PAR_r-2)/2)*PAR_mu0;
  }
  else
  {
	visco = PAR_mu0;
  }
  diverg = U[0]->grad(0) + U[1]->grad(1) + U[2]->grad(2);
  
  delta = pow(U[0]->ip3d->getDJ()/6, 2./3.);
}

CVar C3DNonNewtonianStokesProblem::volumeIntegral(int comp, int elem, int qp)
{
  switch (comp)
  {
    case 0:
    case 1:
    case 2:
      if (tau>0)
      {
        return (U[comp]->val*3-U[comp]->o_val.getVal()*4+U[comp]->oo_val.getVal())/(tau*2)*phi
                +(
                     (U[comp]->grad(0)+U[0]->grad(comp))*gphi(0)
		            +(U[comp]->grad(1)+U[1]->grad(comp))*gphi(1)
		            +(U[comp]->grad(2)+U[2]->grad(comp))*gphi(2)
		         )*0.5*visco.getVal()
                +(U[0]->val*U[comp]->grad(0) + U[1]->val*U[comp]->grad(1) + U[2]->val*U[comp]->grad(2))*phi*conv
                -U[3]->val*gphi(comp) + diverg*gphi(comp)*delta
                -phi*(comp==2?-9.81:0);
      }
      else
      {
        return visco*(
                    (
                     (U[comp]->grad(0)+U[0]->grad(comp))*gphi(0)
		    +(U[comp]->grad(1)+U[1]->grad(comp))*gphi(1)
		    +(U[comp]->grad(2)+U[2]->grad(comp))*gphi(2)
		    )*0.5
		   )
            +(U[0]->val*U[comp]->grad(0) + U[1]->val*U[comp]->grad(1) + U[2]->val*U[comp]->grad(2))*phi*conv
            -U[3]->val*gphi(comp) + diverg*gphi(comp)*delta
            -phi*(comp==2?-9.81:0);
      }
      break;
    case 3:
      return phi*diverg + (U[3]->grad|gphi)*delta;
      break;

  }
  return 0;
}

CVar C3DNonNewtonianStokesProblem::faceIntegral(int comp, int elem)
{
	if (mesh->getFacePhysicalNo(elem) != mortar_pid) return 0;

	double u, v;

	int e = locator->find_cell(x.getVal(), y.getVal(), z.getVal()-GAP, u, v);
//	printf("face %d point (%e, %e, %e) -> face %d\n", elem, x.getVal(), y.getVal(), z.getVal()-GAP, e);

	int offset = 0;

	for (int c=0; c<comp; c++) offset += U[c]->ip3d->getMeshNdof();

	CVar secondary_value = q(offset+U[comp]->ip2d->mapDof(e, 0))*(1e0-u-v)
			              +q(offset+U[comp]->ip2d->mapDof(e, 1))*u
			              +q(offset+U[comp]->ip2d->mapDof(e, 2))*v;

	return (U[comp]->val-secondary_value)*phi*1e6;
}


CVar C3DNonNewtonianStokesProblem::edgeIntegral(int comp, int elem)
{
  return 0;
}

void C3DNonNewtonianStokesProblem::saveOutflow()
{ // save velocity fields along the 'outflow' segments
  for (map<int, bool>:: iterator it = outflows.begin(); it != outflows.end(); it++)
  {
    ofstream f((params["savebase"] + "_" + mesh->getPhysicalName(it->first) + ".dat").c_str());
    for (int e=0; e<mesh->getNfaces(); e++)
    {
      if (mesh->getFacePhysicalNo(e) == it->first)
      {
        for (int i=0; i<3; i++)
	{
	  int ind = mesh->getFaceNode(e, i);
	  f << mesh->getCoord(ind, 0).getVal() << " "
	    << mesh->getCoord(ind, 1).getVal() << " "
	    << mesh->getCoord(ind, 2).getVal() << " "
	    << q(ind).getVal() << " "
	    << q(mesh->getNnodes()+ind).getVal() << " "
	    << q(2*mesh->getNnodes()+ind).getVal() << endl;
	}
      }
    }
    f.close();
  }
}

void C3DNonNewtonianStokesProblem::savePressure()
{ // save pressure along the 'outside' segment
  for (map<int,string>::iterator it=mesh->beginPhysicalNames(); it != mesh->endPhysicalNames(); it++)
  {
    ofstream f((params["savebase"] + "_pressure_" + it->second + ".dat").c_str());
    for (int e=0; e<mesh->getNfaces(); e++)
    {
      if (mesh->getFacePhysicalNo(e) == it->first)
      {
        for (int i=0; i<3; i++)
	{
	  int ind = mesh->getFaceNode(e, i);
	  f << mesh->getCoord(ind, 0).getVal() << " "
	    << mesh->getCoord(ind, 1).getVal() << " "
	    << mesh->getCoord(ind, 2).getVal() << " "
	    << q(3*mesh->getNnodes()+((QuadMesh *)mesh)->getNoOfLnode(ind)).getVal() << endl;
	}
      }
    }
    f.close();
  }
}






int main(int argc, char **argv)
{
/*  if (argc != 2)
  {
    cout << "Usage: " << argv[0] << " problem.conf\n\n";
    exit(1);
  }*/

  PetscInitialize(&argc, &argv, (char *)0, "");

  C3DNonNewtonianStokesProblem P;
  
  P.readConfig(argv[1]);
  
  P.prepare();
  if (P.isTimeDependent())
  {
    P.tSolve(1.0);
  }
  else
  {
    P.nlSolve(1.0);
    P.saveSolution();
    P.saveVTKview();
//    P.calculateToF();
  }

  PetscFinalize();

  return 0;
}
